import os, argparse, yaml, sys
from multiprocessing import Pool

from tools.RAW_to_flatten import changan_flatten
from tools.get_category import get_category
from tools.get_cam_para import get_cam_para
from tools.get_timestamp import get_timestamp
from tools.flatten_to_NuScenes import flatten_nucenes

def multi_process(f, path_list: list):
    pool = Pool(20)
    res = pool.imap(f, path_list)
    for i in res:
        print(i)
        sys.stdout.flush()

def changan_bev_converter(args):
    
    print('Start Converting!')
    
    data_path = args.data_path
    clip_num = args.clip_num
    save_path = args.save_path
    
    # prepare dir
    if not os.path.exists(save_path):
        os.makedirs(save_path, exist_ok=True)
        
    # convert to flatten
    changan_flatten(data_path, clip_num, save_path)
    
    # Prepare configurations
    config_dict = {}
    config_dict['DATA_ROOT_PATH'] = os.path.join(save_path, 'changan_sample_data')
    config_dict['FLATTEN_JSON_ROOT_PATH'] = os.path.join(save_path, 'flatten_json')
    config_dict['OUTPUT_ROOT_PATH'] = os.path.join(save_path, 'nuscenes')
    
    config_dict['NUMBER_OF_SAMPLES'] = 200 * clip_num
    config_dict['DESCRIPTIPTION'] = ''
    config_dict['VEHICLE'] = 'n01'
    config_dict['DATE_CAPTURED'] = '2023-12-09'
    config_dict['LOCATION'] = 'CHINA'
    
    config_dict['CATEGORY_LIST'] = get_category(os.path.join(save_path, 'changan_sample_data'))
    config_dict['SENSOR_DICT'] = {
        'CAM_FRONT': 'camera',
        'CAM_FRONT_LEFT': 'camera',
        'CAM_FRONT_RIGHT': 'camera',
        'CAM_BACK': 'camera',
        'CAM_BACK_LEFT': 'camera',
        'CAM_BACK_RIGHT': 'camera',
        'LIDAR_TOP': 'lidar',
    }
    config_dict['IMG_SUFFIX'] = 'jpeg'
    config_dict['IMG_SIZE'] = {
        'CAM_BACK': [1536, 1920],
        'CAM_BACK_LEFT': [1536, 1920],
        'CAM_BACK_RIGHT': [1536, 1920],
        'CAM_FRONT': [2160, 3840],
        'CAM_FRONT_LEFT': [1536, 1920],
        'CAM_FRONT_RIGHT': [1536, 1920],
    }
    
    config_dict['CALIBRATED_SENSOR'] = get_cam_para('param')
    config_dict['TIMESTAMP_DICT'] = get_timestamp(os.path.join(save_path, 'flatten_json'))
    
    CONFIG_PATH = os.path.join(save_path, 'flatten.yaml')
    with open(CONFIG_PATH, 'w') as f:
        yaml.safe_dump(config_dict, f)
    
    # convert flatten to nuScenes format
    flatten_nucenes(CONFIG_PATH)
    
    print(f'Finish Converting {clip_num} clip!')



if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_path', type=str, required=True, help="Data path to changan bev")
    parser.add_argument('--clip_num', type=int, required=True, help="Clip num you want to convert")
    parser.add_argument('--save_path', type=str, required=True, help="Save path to nuScenes format data")
    args = parser.parse_args()
    
    changan_bev_converter(args)